package com.lzqinga.service;

import com.google.gson.JsonObject;
import com.google.gson.JsonParser;
import okhttp3.*;
import org.apache.commons.codec.binary.Base64;
import org.json.JSONObject;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.concurrent.TimeUnit;

/**
 * - @author lzqinga
 * - @description
 * - @date 2025/10/10/010 15:28
 **/
public class ImagesService {
    public static final OkHttpClient HTTP_CLIENT = new OkHttpClient().newBuilder()
            .readTimeout(300, TimeUnit.SECONDS)
            .build();

    // 默认参数配置
    private static final String DEFAULT_SIZE = "1024x768";
    private static final int DEFAULT_N = 1;
    private static final int DEFAULT_STEPS = 20;
    private static final String DEFAULT_SAMPLER = "Euler a";
    private static final String DEFAULT_USER_ID = "123";

    private static final String API_TOKEN = "bce-v3/ALTAK-7E9eb6lRRPwixxdqYc76p/786cadc5603df42c8459dc24aea8df3d175b6548";

    /**
     * 生成图片的封装方法
     * @param prompt 正面提示词
     * @param negativePrompt 负面提示词
     * @param size 图片尺寸
     * @param n 生成数量
     * @param steps 步数
     * @param sampler 采样器
     * @param userId 用户ID
     * @param token 认证token
     * @return API响应结果
     * @throws IOException 网络请求异常
     */
    public static String generateImage(String prompt, String negativePrompt, String size,
                                       Integer n, Integer steps, String sampler,
                                       String userId, String token) throws IOException {
        // 构建请求体JSON
        String jsonBody = String.format(
                "{\"prompt\":\"%s\",\"negative_prompt\":\"%s\",\"size\":\"%s\",\"n\":%d,\"steps\":%d,\"sampler_index\":\"%s\",\"user_id\":\"%s\"}",
                prompt, negativePrompt, size, n, steps, sampler, userId
        );

        MediaType mediaType = MediaType.parse("application/json");
        RequestBody body = RequestBody.create(mediaType, jsonBody);

        Request request = new Request.Builder()
                .url("https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/text2image/sd_xl")
                .method("POST", body)
                .addHeader("Content-Type", "application/json")
                .addHeader("Accept", "application/json")
                .addHeader("Authorization", "Bearer " + token)
                .build();

        try (Response response = HTTP_CLIENT.newCall(request).execute()) {
            return response.body().string();
        }
    }

    /**
     * 简化版生成方法（使用默认参数）
     * @param prompt 正面提示词
     * @param negativePrompt 负面提示词
     * @param token 认证token
     * @return API响应结果
     * @throws IOException 网络请求异常
     */
    public static String generateImage(String prompt, String negativePrompt, int num, String token) throws IOException {
        return generateImage(prompt, negativePrompt, DEFAULT_SIZE, num,
                DEFAULT_STEPS, DEFAULT_SAMPLER, DEFAULT_USER_ID, token);
    }

    /**
     * 最简版生成方法（只需提示词和token）
     * @param prompt 正面提示词
     * @param num 生成图片数量
     * @param token 认证token
     * @return API响应结果
     * @throws IOException 网络请求异常
     */
    public static String generateImage(String prompt, int num, String token) throws IOException {
        return generateImage(prompt, "", num,  token);
    }

    public static String saveImage(String jsonResponse, String savePath) throws IOException {
        // 解析JSON响应
        JsonObject jsonObject = JsonParser.parseString(jsonResponse).getAsJsonObject();

        if (jsonObject.has("data") && jsonObject.get("data").isJsonArray()) {
            JsonObject data = jsonObject.get("data").getAsJsonArray().get(0).getAsJsonObject();

            if (data.has("b64_image")) {
                String base64Image = data.get("b64_image").getAsString();

                // 移除Base64前缀（如果有）
                if (base64Image.contains(",")) {
                    base64Image = base64Image.split(",")[1];
                }

                // 解码Base64并保存图片
                byte[] imageBytes = Base64.decodeBase64(base64Image);

                File outputFile = new File(savePath);

                // 确保目录存在
                outputFile.getParentFile().mkdirs();

                // 保存图片
                try (FileOutputStream fos = new FileOutputStream(outputFile)) {
                    fos.write(imageBytes);
                }

                return outputFile.getAbsolutePath();
            }
        }

        throw new IOException("未找到图片数据");
    }

    public static void main(String[] args) throws IOException {
        String url = generateImage("美丽的女孩在乡间的小路上，粒子组成的蝴蝶飞舞，侧脸，近景", 1, API_TOKEN);
        saveImage(url, "output/videos");
        System.out.println("url = " + url);
    }
}
